from rex.exploit.cgc import CGCType2Exploit
from ..c_templates import c_template_type2_circumstantial

# todo equation solving
# todo randomness
class CGCType2GeneralExploit(CGCType2Exploit):
    """
    A CGC type2 exploit object, which leaks memory from the flag page
    Might be generated circumstantially or through another technique
    """

    def __init__(self, crash, input_str, output_index, bypasses_nx, bypasses_aslr, method_name='unclassified'):
        """
        :param crash: a crash object which has been modified to exploit a vulnerability
        :param addr_var: the claripy variable used to determine the position in the payload which contains the address
        """
        super(CGCType2GeneralExploit, self).__init__(crash, bypasses_nx=bypasses_nx, bypasses_aslr=bypasses_aslr)

        self.method_name = method_name
        self.output_leak_idx = output_index
        self._raw_payload = input_str

    def __str__(self):
        return "general type2"

    def test(self):
        raise NotImplementedError

    def dump_c(self, filename=None):
        """
        Creates a simple C file to do the type1 exploit
        :param filename: dumps the code to this path if filename is not None
        :return: the c_code
        """
        encoded_payload = ""
        for c in self._raw_payload:
            encoded_payload += "\\x%02x" % c

        fmt_args = dict()
        fmt_args["payload"] = encoded_payload
        fmt_args["payloadsize"] = str(len(self._raw_payload))
        fmt_args["output_leak_idx"] = str(self.output_leak_idx)

        # TODO using .format is annoying because of all the curly braces
        # figure out how to do this better
        c_code = c_template_type2_circumstantial.c_template
        for k, v in fmt_args.items():
            c_code = c_code.replace("{%s}" % k, v)

        if filename is not None:
            with open(filename, 'w') as f:
                f.write(c_code)
        else:
            return c_code
